import time

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import pandas as pd
import math
import cmath
import csv
from scipy.fftpack import fftshift, fftn
from matplotlib.colors import BoundaryNorm
import cmocean

# 光速
c = 3e8  # m/s
# 角度制转换
def degree_to_rad(d):
    return d/180*math.pi
# 两点距离
def dist(a,b):
    return math.sqrt((a[0]-b[0])**2+(a[1]-b[1])**2+(a[2]-b[2])**2)

# 分贝转换
def k_dB(xt, r1020 = 20):
    if np.ndim(xt) == 0:
        return r1020*math.log10(xt)
    else:
        ann = np.zeros(len(xt))
        for i in range(len(xt)):
            ann[i] = r1020*math.log10(xt[i])
        return ann

def dB_k(xt, r1020 = 20):
    if np.ndim(xt)==0:
        return 10**(xt/r1020)
    else:
        ann = np.zeros(len(xt))
        for i in range(len(xt)):
            ann[i] = 10**(xt[i]/r1020)
        return ann

# 二维区间积分器
def jifenqi_2d(x1,x2,y1,y2,f, rk, x_s, dx, showit = False):
    jifendd = []
    df =dx
    summ = 0
    si = (x2-x1)/df * (y2-y1)/df
    xi = np.linspace(start=x1+x_s, stop=x2+x_s, num=df) + (x2-x1)/df/2
    xi = xi[:-1]
    yi = np.linspace(start=y1, stop=y2, num=df) + (y2-y1)/df/2
    yi = yi[:-1]
    for i in xi:
        for j in yi:
            if rk > 0:
                if (i-x_s)**2 + j*j < rk**2:
                    summ = summ + f(i, j)*si
                    jifendd.append([i, j])
            else:
                summ = summ + f(i, j) * si
                jifendd.append([i, j])

    if showit:
        jifendd = np.array(jifendd)
        plt.scatter(jifendd[:, 1], jifendd[:,0])
        plt.axis('equal')
        plt.grid()
        plt.show()
    return summ


class Reflect_Array():

    def __init__(self, f, R_base, h, acell, RP_LABA, RP_ele, T):
        self.f = f  # Hz, 1/s
        self.lam_0 = c/f # m
        self.k_0 = 2 * math.pi * f / c  # rad/m
        self.R_base = R_base
        self.h = h
        self.T = T
        self.acell = acell

        self.sbili = self.acell/self.lam_0 * 2
        self.bians = int(self.T * (1 - self.sbili) * 0.5)
        # print("bians: ", self.sbili, self.bians)

        self.fxtMask = np.zeros((self.T, self.T))
        self.fxt_ww = np.zeros((self.T, self.T))

        for i in range(self.T):
            for j in range(self.T):
                u = i / (self.T - 1) * 2 - 1
                v = j / (self.T - 1) * 2 - 1
                u = u / self.sbili
                v = v / self.sbili
                if u * u + v * v <= 1:
                    self.fxtMask[i, j] = 1
                    self.fxt_ww[i, j] = np.sqrt(abs(1 - (u * u + v * v)))

        # self.plot_ELE(self.fxtMask)
        # self.plot_ELE(self.fxt_ww)
        self.a_all = R_base * R_base * math.pi
        self.n = int(self.R_base * 2 / self.acell)
        self.s_ele = acell * acell

        self.laba_in = [0, 0, h]
        self.r_0 = dist(self.laba_in, [0, 0, 0])

        self.RP_LABA = RP_LABA
        self.RP_ele = RP_ele
        self.if_ideal = False

        xx = np.arange(-self.R_base, self.R_base, acell) + acell / 2
        yy = np.arange(-self.R_base, self.R_base, acell) + acell / 2

        self.rmn = np.zeros((self.n, self.n))
        self.arv = np.zeros((self.n, self.n))
        self.pos_mn = np.zeros((self.n, self.n, 2))
        self.pha_mn = np.zeros((self.n, self.n))
        self.mag_mn = np.ones((self.n, self.n))
        self.daX = np.zeros((self.T, self.T), dtype=complex)

        self.phhh_const = 0
        pp = np.linspace(0, math.pi * 2, 17)
        sum_fenmu = 0
        sum_fenzi = 0

        for i in range(self.n):
            for j in range(self.n):
                elepos = [xx[i], yy[j]]
                self.pos_mn[i, j] = elepos
                self.rmn[i, j] = math.sqrt(elepos[0] ** 2 + elepos[1] ** 2 + self.h ** 2)
                smn = math.sqrt(elepos[0] ** 2 + elepos[1] ** 2)
                thhh = math.acos((self.rmn[i, j] ** 2 + self.r_0 ** 2 - smn ** 2) / (2 * self.rmn[i, j] * self.r_0))
                thhh2 = math.acos(self.h / self.rmn[i, j])

                if np.sqrt(elepos[0] ** 2 + elepos[1] ** 2) < self.R_base - self.acell * 0.6:
                    self.arv[i, j] = self.RP_LABA(thhh, self.phhh_const) * self.RP_ele(thhh2, 0) / self.rmn[i, j]
                else:
                    self.arv[i, j] = 0

                ssa = 0 # self.k_0 * 0.5 * elepos[0]
                self.pha_mn[i, j] = (self.k_0 * self.rmn[i, j] + ssa) % (math.pi * 2)
                # self.pha_mn[i, j] = pp[np.argmin(abs(self.pha_mn[i, j] - pp))]
                self.daX[i, j] = self.mag_mn[i, j] * self.arv[i, j] * cmath.exp(1j*(self.pha_mn[i, j] - self.k_0 * self.rmn[i, j]))
                sum_fenzi = sum_fenzi + self.daX[i, j] * self.s_ele
                sum_fenmu = sum_fenmu + self.mag_mn[i, j] * self.arv[i, j] * self.mag_mn[i, j] * self.arv[i, j] * self.s_ele
        # u_ill = (sum_fenzi * sum_fenzi)/(sum_fenmu * self.a_all)

    def Change_ele(self, pha_mn_d, mag_mn_d):
        for i in range(self.n):
            for j in range(self.n):
                self.pha_mn[i, j] = pha_mn_d[i, j].real
                self.mag_mn[i, j] = mag_mn_d[i, j].real
                # self.daX[i, j] = self.mag_mn[i, j] * self.s_ele * self.arv[i, j] * cmath.exp(1j * (self.pha_mn[i, j] - self.k_0 * self.rmn[i, j]))

    # 单元辐射强度分布
    def Am_mn(self, elepos):
        rmn = dist(elepos, self.laba_in)
        smn = dist(elepos, [0, 0, 0])
        thhh = math.acos((rmn**2 + self.r_0**2 - smn**2)/(2 * rmn * self.r_0))
        thhh2 = math.acos(self.h / rmn)
        A_result = self.RP_LABA(thhh, self.phhh_const) * self.RP_ele(thhh2, 0) / rmn
        return A_result, thhh2

    def plot_ELE(self, q_ele):
        uq_ele = q_ele.copy()
        uq_ele = (uq_ele * 360 / math.pi) % (360)
        # print(uq_ele)
        uq_ele[uq_ele < 0.001] = -10
        plt.title("Phase Distribution")
        lvls = np.linspace(0, 360, 13)

        cmap1 = plt.get_cmap('twilight')
        norm1 = BoundaryNorm(lvls, cmap1.N)

        imm = plt.imshow(uq_ele, norm=norm1, cmap=cmap1)
        imm = plt.imshow(uq_ele)
        plt.colorbar(imm)
        plt.show()
        # X = np.arange(0, self.n, step=1)  # X轴的坐标
        # Y = np.arange(0, self.n, step=1)  # Y轴的坐标
        # # 设置每一个（X，Y）坐标所对应的Z轴的值，在这边Z（X，Y）=X+Y
        #
        # xx, yy = np.meshgrid(X, Y)  # 网格化坐标
        # X, Y = xx.ravel(), yy.ravel()  # 矩阵扁平化
        # bottom = np.zeros_like(X)  # 设置柱状图的底端位值
        # Z = (q_ele.ravel() % (math.pi * 2)) / 500  # 扁平化矩阵
        #
        #
        # ax = plt.axes(projection='3d')  # 三维坐标轴
        # ax.bar3d(X, Y, bottom, 1, 1, Z, shade=True)  #
        # # 坐标轴设置
        # plt.show()

    def plot_FFT(self):
        Mc = len(self.fxt)
        x = np.linspace(-1, 1, Mc)
        y = np.linspace(-1, 1, Mc)
        X, Y = np.meshgrid(x, y)

        # ax3 = plt.axes(projection='3d')
        # ax3.plot_surface(Y, X, self.fxt, cmap='rainbow')
        # plt.show()

        # plt.subplot(1, 2, 1)
        # cset = plt.contourf(X, Y, self.fxt, 10, cmap=plt.cm.hot)
        # contour = plt.contour(X, Y, self.fxt, 10, colors='k')
        # plt.clabel(contour, fontsize=10, colors='k')
        # plt.axis('equal')
        # plt.colorbar(cset)
        #
        # plt.subplot(1, 2, 2)
        fxtt = self.fxt.copy()
        # fxtt[fxtt > 0] = fxtt[fxtt > 0] + 10**(-3.5)
        lg_fxt = 10 * np.log10(fxtt)
        lvls = np.array([-35, -30, -25, -20, -15, -10, -5, -3, 0])
        # lvls = np.arange(-36, 0.1, 2)
        # cset = plt.contourf(X, Y, lg_fxt, lvls, cmap="rainbow",)
        cset = plt.contourf(Y, X, lg_fxt, lvls, cmap=mpl.cm.jet, )
        # contour = plt.contour(X, Y, lg_fxt, lvls, colors='k', )
        # plt.clabel(contour, fontsize=10, colors='k')

        plt.title("Radiation Pattern")
        plt.axis('equal')
        plt.colorbar(cset)
        plt.show()

    def FFT_to_FXT(self):
        self.sx = fftn(self.daX)
        # print(np.shape(self.daX))
        self.fft_sx = fftshift(self.sx)
        self.fxt_ = np.abs(self.fft_sx) ** 2
        self.fxt_ = self.fxtMask * self.fxt_ww * self.fxt_
        self.fxt = self.fxt_ / np.max(self.fxt_)
        # self.fxt = self.fxt_[self.bians:self.T-self.bians, self.bians:self.T-self.bians]

if __name__ == '__main__':
    h = 0.11
    f0 = 30e9
    lam0 = c/f0
    L = 9 * lam0
    k_0 = 2 * math.pi * f0 / c  # rad/m
    aa = lam0 * 0.5
    from RA_Cal import Para_RA
    TA1 = Reflect_Array(f=f0, R_base=L, h=h, RP_LABA=Para_RA.RP_LABA, RP_ele=Para_RA.RP_ele, acell=aa, T=1001)
    # TA1.plot_ELE(TA1.pha_mn)
    # TA1.plot_ELE(np.abs(TA1.daX)[:TA1.n, :TA1.n])
    # TA1.plot_ELE(np.angle(TA1.daX)[:TA1.n, :TA1.n])
    TA1.FFT_to_FXT()
    TA1.plot_FFT()

    Mc = len(TA1.fxt)
    M_2 = int(Mc/2)
    xx = np.linspace(-1, 1, Mc)
    # xx = np.arcsin(xx) / math.pi * 180
    print(np.argmax(TA1.fxt))
    plt.plot(xx, list(TA1.fxt[:, M_2]))
    plt.grid()
    plt.show()

    db_fxt = np.log10(TA1.fxt[:, M_2]) * 10
    plt.plot(xx, list(db_fxt))
    plt.grid()
    plt.show()

    # with open("FFT_Result_Zmith.csv", 'a', encoding='utf-8', newline='') as f:
    #     write = csv.writer(f)  # 创建writer对象
    #     write.writerow(list(xx))
    #     write.writerow(list(db_fxt))
